import cv2
import numpy as np
import glob
from tqdm import tqdm

# 定义棋盘大小: 注意此处是内部的行、列角点个数，不包含最外边两列，否则会出错
chessboard_size = (7, 6)

# 生成54×3的矩阵，用来保存棋盘图中9*6个内角点的3D坐标，也就是物体点坐标
objp = np.zeros((np.prod(chessboard_size), 3), dtype=np.float32)
# 通过np.mgrid生成对象的xy坐标点，每个棋盘格大小是130mm
# 最终得到z=0的objp为(0,0,0), (1*13,0,0), (2*13,0,0) ,...
objp[:, :2] = np.mgrid[0:chessboard_size[0], 0:chessboard_size[1]].T.reshape(-1, 2) * 13

# print("object is %f", objp)

# 定义数组，来保存监测到的点
obj_points = []  # 保存世界坐标系的三维点
img_points = []  # 保存图片坐标系的二维点

# 设置终止条件： 迭代30次或者变动 < 0.001
criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)

# 读取目录下的所有图片
calibration_paths = glob.glob('./img*.jpg')

# 为方便显示，使用tqdm显示进度条
for image_path in tqdm(calibration_paths):
    # 读取图片
    img = cv2.imread(image_path)
    # 图像二值化
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    # 找到棋盘格内角点位置
    ret, corners = cv2.findChessboardCorners(gray, chessboard_size, None)

    if ret == True:
        obj_points.append(objp)
        # 亚像素级角点检测，在角点检测中精确化角点位置
        corners2 = cv2.cornerSubPix(gray, corners, (11, 11), (-1, -1), criteria)
        img_points.append(corners2)

        # 在图中标注角点,方便查看结果
        img = cv2.drawChessboardCorners(img, chessboard_size, corners2, ret)
        cv2.namedWindow('img', 0)
        cv2.resizeWindow('img', 900, 600)
        cv2.imshow('img', img)
        cv2.waitKey(500)
cv2.destroyAllWindows()
print("ok")
# 相机标定
ret, mtx, dist, rvecs, tvecs = cv2.calibrateCamera(obj_points, img_points, gray.shape, None, None)

# 保存相机参数
np.savez('C.npz', mtx=mtx, dist=dist, rvecs=rvecs, tvecs=tvecs)
print("ret:", ret)
print("内参数矩阵:\n", mtx)
print("畸变系数:\n", dist)
print("旋转向量:", rvecs)  # 外参数
print("平移向量:", tvecs)  # 外参数
print("相机标定结束")


with np.load('C.npz') as X:
    mtx, dist, _, _ = [X[i] for i in ('mtx', 'dist', 'rvecs', 'tvecs')]

# 像素坐标
test_img = cv2.imread("./1m2.jpg")
gray = cv2.cvtColor(test_img, cv2.COLOR_BGR2GRAY)


# 找到图像平面点角点坐标
ret, corners = cv2.findChessboardCorners(gray, chessboard_size, None)

retval, rvec, tvec = cv2.solvePnP(objp, corners, mtx, dist)
print("旋转矩阵", rvec)
print("平移变量", tvec)
